#!/usr/bin/env python

#
# This file contains functions for dealing with the blimp wireless
# communication protocol.
#

import struct
import serial
import time
import numpy as np

FRAME_SIZE = 20
DATA_SIZE  = FRAME_SIZE - 4
BAUD_RATE  = 115200

class States:
    IDLE = 0
    B1   = 1
    DATA = 2
    E1   = 3
    SEND = 4

class Datalink(object):
    SS = chr(0x55)
    ES = chr(0xAA)

    def __init__(self, port, baudrate = BAUD_RATE):
        self.s = serial.Serial(port=port, baudrate=baudrate)
        self.state = States.IDLE

    def create_message(self, opcode, data):
        # opcode is a single byte
        # data is a tuple
        pass

    def tx_raw(self, msg):
        msg_len = len(msg)
        if msg_len != FRAME_SIZE:
            raise ValueError("msg was %d bytes, it must be %d"%(msg_len,
                                                                DATA_SIZE))
        self.s.write(msg)
        return

    def rx_raw(self, bytes = FRAME_SIZE):
        return self.s.read(bytes)

    def receive(self):
        done = False
        num_data = 0
        data_idx = 0
        data = [0 for i in range(DATA_SIZE)]
        while not done:
            byte = self.s.read(1)
            if self.state == States.IDLE:
                if byte == self.SS:
                    self.state = States.B1
                else:
                    self.state = States.IDLE
                    
            elif self.state == States.B1:
                if byte == self.SS:
                    self.state = States.DATA
                    data_idx = 0
                else:
                    self.state = States.IDLE
                    
            elif self.state == States.DATA:
                if data_idx == DATA_SIZE:
                    if byte == self.ES:
                        self.state = States.E1
                    else:
                        self.state = States.IDLE
                else:
                    data[data_idx] = byte
                    data_idx += 1

            elif self.state == States.E1:
                if byte == self.ES:
                    done = True
                    self.state = States.IDLE
                else:
                    self.state = States.IDLE
        return data

    def send(self, data):
        # data - 16 chars in an array or string
        data = list(data)
        if len(data) < 16:
            data += [chr(0)] * (16 - len(data))
        self.tx_raw(self.SS * 2 + ''.join(data) + self.ES * 2)


class ExpFilter(object):
    # Exponential decay LPF
    def __init__(self, alpha=1/4.):
        self.alpha = alpha   

    def __call__(self, new_data):
        try:
            self.data
        except AttributeError:
            self.data = new_data

        self.data = self.alpha*new_data + (1 - self.alpha)*self.data
        return self.data

class SensorInfo(object):
    channels = ['front', 'reverse', 'left', 'right', 'down']
    
    def __init__(self, datalink, filter_class=ExpFilter, args=()):
        self.dl = datalink
        if args:
            self.filter = filter_class(*args)
        else:
            self.filter = filter_class()
    
    def get_sensors(self):
        data = np.zeros(5)              # 5 channels of data
        
        now = time.time()           # For measuring control loop rate
        new_data = np.array(struct.unpack("hhhhhh", ''.join(self.dl.receive())[:12])[1:])

        # Apply our filter
        data = self.filter(new_data)

        # Calculate the refresh rate
        rate = 1.0/(time.time() - now)

        # Create a dictionary indexed by channel name
        ir = dict(zip(self.channels, data))

        # linearize the data
        down = ir['down']
        ir['down'] = 0.5821*down**2 - 35.163*down + 802.22
        #ir['rate'] = rate

        return ir

class MotorControl(object):
    def __init__(self, datalink):
        self.dl = datalink

    def set_motors(self, l,r,v):
        self.dl.send(struct.pack("bbbb", 1, l,r,v))

        
            
